package core

import (
	"crypto/rand"
	"encoding/hex"
	"fmt"
	"log"
	"net"
	"strings"
	"sync"
	"time"

	"dnslog/config"

	"github.com/miekg/dns"
)

// DNSRecord 表示一条DNS查询记录
type DNSRecord struct {
	ID        string    `json:"id"`
	Subdomain string    `json:"subdomain"`
	SourceIP  string    `json:"source_ip"`
	Timestamp time.Time `json:"timestamp"`
	QueryType string    `json:"query_type"`
}

// DNSLogPlatform DNSLOG平台
type DNSLogPlatform struct {
	config      *config.Config
	domain      string
	records     map[string][]DNSRecord
	recordsLock sync.RWMutex
	users       map[string]string // user_token -> subdomain
	usersLock   sync.RWMutex
}

// NewDNSLogPlatform 创建新的DNSLOG平台
func NewDNSLogPlatform(cfg *config.Config) *DNSLogPlatform {
	platform := &DNSLogPlatform{
		config:  cfg,
		domain:  cfg.Server.Domain,
		records: make(map[string][]DNSRecord),
		users:   make(map[string]string),
	}

	// 启动记录清理任务
	go platform.startCleanupTask()

	return platform
}

// GenerateRandomToken 生成随机用户token
func (p *DNSLogPlatform) GenerateRandomToken() (string, error) {
	bytes := make([]byte, 16)
	if _, err := rand.Read(bytes); err != nil {
		return "", err
	}
	return hex.EncodeToString(bytes), nil
}

// GenerateSubdomain 为用户生成基础子域名
func (p *DNSLogPlatform) GenerateSubdomain(userToken string) string {
	randomPart, _ := p.GenerateRandomToken()
	baseSubdomain := randomPart[:8] + "." + p.domain

	p.usersLock.Lock()
	p.users[userToken] = baseSubdomain
	p.usersLock.Unlock()

	log.Printf("🎯 Generated base subdomain: %s for token: %s", baseSubdomain, userToken)
	return baseSubdomain
}

// GetUserSubdomain 获取用户的基础子域名
func (p *DNSLogPlatform) GetUserSubdomain(userToken string) string {
	p.usersLock.RLock()
	defer p.usersLock.RUnlock()
	return p.users[userToken]
}

// extractBaseSubdomain 提取基础子域名
func (p *DNSLogPlatform) extractBaseSubdomain(fullDomain string) string {
	// 去掉末尾的点
	domain := fullDomain
	if len(domain) > 0 && domain[len(domain)-1] == '.' {
		domain = domain[:len(domain)-1]
	}

	// 找到基础子域名部分（随机部分 + 主域名）
	// 例如：Administrator.558b7e7d.aaa.ccc -> 558b7e7d.aaa.ccc
	parts := strings.Split(domain, ".")
	if len(parts) < 3 {
		return domain // 如果没有前缀，返回原域名
	}

	// 取最后三部分作为基础子域名
	baseParts := parts[len(parts)-3:]
	baseSubdomain := strings.Join(baseParts, ".")

	return baseSubdomain
}

// AddDNSRecord 添加DNS记录
func (p *DNSLogPlatform) AddDNSRecord(fullDomain, sourceIP, queryType string) {
	// 提取基础子域名
	baseSubdomain := p.extractBaseSubdomain(fullDomain)
	if baseSubdomain == "" {
		log.Printf("❌ Cannot extract base subdomain from: %s", fullDomain)
		return
	}

	record := DNSRecord{
		ID:        fmt.Sprintf("%d", time.Now().UnixNano()),
		Subdomain: fullDomain, // 保存完整的域名
		SourceIP:  sourceIP,
		Timestamp: time.Now(),
		QueryType: queryType,
	}

	p.recordsLock.Lock()
	defer p.recordsLock.Unlock()

	// 检查记录数量限制
	if len(p.records[baseSubdomain]) >= p.config.Security.MaxRecordsPerUser {
		// 移除最旧的记录
		p.records[baseSubdomain] = p.records[baseSubdomain][1:]
	}

	// 使用基础子域名作为key
	p.records[baseSubdomain] = append(p.records[baseSubdomain], record)

	log.Printf("💾 Saved DNS record: %s (base: %s) from %s, type: %s",
		fullDomain, baseSubdomain, sourceIP, queryType)
}

// GetDNSRecords 获取用户的DNS记录
func (p *DNSLogPlatform) GetDNSRecords(userToken string) []DNSRecord {
	baseSubdomain := p.GetUserSubdomain(userToken)
	if baseSubdomain == "" {
		log.Printf("❌ GetDNSRecords: No subdomain found for token: %s", userToken)
		return []DNSRecord{}
	}

	p.recordsLock.RLock()
	defer p.recordsLock.RUnlock()

	records := p.records[baseSubdomain]
	log.Printf("🔍 GetDNSRecords: Found %d records for base subdomain: %s", len(records), baseSubdomain)

	// 返回记录的副本，避免并发问题
	result := make([]DNSRecord, len(records))
	copy(result, records)
	return result
}

// ClearDNSRecords 清空用户的DNS记录
func (p *DNSLogPlatform) ClearDNSRecords(userToken string) {
	baseSubdomain := p.GetUserSubdomain(userToken)
	if baseSubdomain == "" {
		return
	}

	p.recordsLock.Lock()
	defer p.recordsLock.Unlock()

	p.records[baseSubdomain] = []DNSRecord{}
	log.Printf("🗑️ Cleared all records for: %s", baseSubdomain)
}

// StartDNSServer 启动DNS服务器
func (p *DNSLogPlatform) StartDNSServer(listenAddr string) error {
	dns.HandleFunc(p.domain+".", p.handleDNSQuery)

	server := &dns.Server{
		Addr: listenAddr,
		Net:  "udp",
	}

	log.Printf("DNS server starting on %s", listenAddr)
	return server.ListenAndServe()
}

// handleDNSQuery 处理DNS查询
func (p *DNSLogPlatform) handleDNSQuery(w dns.ResponseWriter, r *dns.Msg) {
	m := new(dns.Msg)
	m.SetReply(r)
	m.Authoritative = true

	for _, question := range r.Question {
		domain := question.Name
		queryType := dns.TypeToString[question.Qtype]

		// 去掉域名末尾的点
		if len(domain) > 0 && domain[len(domain)-1] == '.' {
			domain = domain[:len(domain)-1]
		}

		// 记录DNS查询
		sourceIP, _, _ := net.SplitHostPort(w.RemoteAddr().String())
		p.AddDNSRecord(domain, sourceIP, queryType)

		// 返回一个A记录响应
		if question.Qtype == dns.TypeA {
			rr := &dns.A{
				Hdr: dns.RR_Header{
					Name:   question.Name,
					Rrtype: dns.TypeA,
					Class:  dns.ClassINET,
					Ttl:    60,
				},
				A: net.ParseIP("127.0.0.1"),
			}
			m.Answer = append(m.Answer, rr)
		}

		// 对于其他查询类型，也返回适当的响应
		if question.Qtype == dns.TypeAAAA {
			rr := &dns.AAAA{
				Hdr: dns.RR_Header{
					Name:   question.Name,
					Rrtype: dns.TypeAAAA,
					Class:  dns.ClassINET,
					Ttl:    60,
				},
				AAAA: net.ParseIP("::1"),
			}
			m.Answer = append(m.Answer, rr)
		}
	}

	w.WriteMsg(m)
}

// GetDomain 获取域名
func (p *DNSLogPlatform) GetDomain() string {
	return p.domain
}

// 启动定期清理任务
func (p *DNSLogPlatform) startCleanupTask() {
	ticker := time.NewTicker(time.Duration(p.config.Security.CleanupInterval) * time.Second)
	defer ticker.Stop()

	for range ticker.C {
		p.cleanupOldRecords()
	}
}

// 清理过期记录
func (p *DNSLogPlatform) cleanupOldRecords() {
	p.recordsLock.Lock()
	defer p.recordsLock.Unlock()

	expirationTime := time.Now().Add(-p.config.GetRecordRetentionDuration())
	cleanedCount := 0

	for subdomain, records := range p.records {
		var validRecords []DNSRecord
		for _, record := range records {
			if record.Timestamp.After(expirationTime) {
				validRecords = append(validRecords, record)
			} else {
				cleanedCount++
			}
		}
		p.records[subdomain] = validRecords

		// 如果该子域名没有记录了，删除空数组
		if len(validRecords) == 0 {
			delete(p.records, subdomain)
		}
	}

	if cleanedCount > 0 {
		log.Printf("🧹 Cleaned %d expired records", cleanedCount)
	}
}
